#include <cuda_runtime.h>
#include <device_launch_parameters.h>

#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION
#define STB_DXT_IMPLEMENTATION
#define STB_TRUETYPE_IMPLEMENTATION
#define LIBRARYNAME_IMPLEMENTATION

#include "stb_image/stb_image.h"
#include "stb_image/stb_image_write.h"

#include "lastError.h"
#include "../cvtColor.h"
#include "../imageFile.h"

namespace Image2D
{

int imageReadToGray      (std::string fileName, Image2D_Ptr& i2d)
{
    int width;
    int height;
    int nChannels;
    unsigned char *data = stbi_load(fileName.c_str(), &width, &height, &nChannels, 1);
    if (data == nullptr)
    {
        LastError::setLastError("error: imageReadToGray stbi_load", -1);
        return LastError::getLastErrorCode();
    }

    i2d = Image2D_Ptr(width, height, ImageFormat::IMAGE_NONE, ImageType::IMAGE_UNSIGNED_CHAR);

    cudaMemcpy2D((void*)(i2d.data(0)), i2d.linesize(0), data, width, width, height,     cudaMemcpyHostToDevice);
    LastError::setLastError("error: imageReadToGray cudaMemcpy2D");

    stbi_image_free(data);

    return LastError::getLastErrorCode();
}

int imageReadToRgb  (std::string fileName, Image2D_Ptr& i2d)
{
    int width;
    int height;
    int nChannels;
    unsigned char *data = stbi_load(fileName.c_str(), &width, &height, &nChannels, 3);
    if (data == nullptr)
    {
        LastError::setLastError("error: imageReadToRgb stbi_load", -1);
        return LastError::getLastErrorCode();
    }

    i2d = Image2D_Ptr(width, height, ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);

    cudaMemcpy2D((void*)(i2d.data(0)), i2d.linesize(0), data, width * 3, width * 3, height,     cudaMemcpyHostToDevice);
    LastError::setLastError("error: imageReadToRgb cudaMemcpy2D");

    stbi_image_free(data);

    return LastError::getLastErrorCode();
}

int imageReadToBgr      (std::string fileName, Image2D_Ptr& i2d)
{
    Image2D_Ptr i2dRgb;
    int ret = imageReadToRgb(fileName, i2dRgb);
    if (ret != 0)
        return ret;

    i2d = Image2D_Ptr(i2dRgb.width(), i2dRgb.height(), ImageFormat::IMAGE_BGR, ImageType::IMAGE_UNSIGNED_CHAR);

    rgbToBgr(i2dRgb, i2d);

    return LastError::getLastErrorCode();
}

int imageReadToYuv420p  (std::string fileName, Image2D_Ptr& i2d)
{
    Image2D_Ptr i2dRgb;
    int ret = imageReadToRgb(fileName, i2dRgb);
    if (ret != 0)
        return ret;

    i2d = Image2D_Ptr(i2dRgb.width(), i2dRgb.height(), ImageFormat::IMAGE_YUV420p, ImageType::IMAGE_UNSIGNED_CHAR);

    rgbToYuv420p(i2dRgb, i2d);

    return LastError::getLastErrorCode();
}

int imageReadToNv12     (std::string fileName, Image2D_Ptr& i2d)
{
    Image2D_Ptr i2dRgb;
    int ret = imageReadToRgb(fileName, i2dRgb);
    if (ret != 0)
        return ret;

    i2d = Image2D_Ptr(i2dRgb.width(), i2dRgb.height(), ImageFormat::IMAGE_NV12, ImageType::IMAGE_UNSIGNED_CHAR);

    rgbToNv12(i2dRgb, i2d);

    return LastError::getLastErrorCode();
}

int imageWriteJpegFromGray    (std::string fileName, const Image2D_Ptr& i2d, int quality)
{
    if((int)(i2d.type()) != (int)(ImageType::IMAGE_UNSIGNED_CHAR))
    {
        LastError::setLastError("error: Image2D_Ptr type must unsigned char", -1);
        return LastError::getLastErrorCode();
    }

    int width = i2d.width();
    int height = i2d.height();
    unsigned char* data = (unsigned char*)malloc(width * height);
    if (data == nullptr)
    {
        LastError::setLastError("error: imageWriteJpegFromGray malloc", -1);
        return LastError::getLastErrorCode();
    }

    cudaMemcpy2D(data, width, i2d.data(0), i2d.linesize(0), width ,height,     cudaMemcpyDeviceToHost);
    LastError::setLastError("error: imageWriteJpegFromGray cudaMemcpy2D");

    stbi_write_jpg(fileName.c_str(), width, height, 1, data, quality);

    free(data);

    return LastError::getLastErrorCode();
}

int imageWriteJpegFromRgb     (std::string fileName, const Image2D_Ptr& i2d, int quality)
{
    if((int)(i2d.type()) != (int)(ImageType::IMAGE_UNSIGNED_CHAR))
    {
        LastError::setLastError("error: Image2D_Ptr type must unsigned char", -1);
        return LastError::getLastErrorCode();
    }

    int width = i2d.width();
    int height = i2d.height();
    unsigned char* data = (unsigned char*)malloc(width * height * 3);
    if (data == nullptr)
    {
        LastError::setLastError("error: imageWriteJpegFromRgb malloc", -1);
        return LastError::getLastErrorCode();
    }

    cudaMemcpy2D(data, width * 3, i2d.data(0), i2d.linesize(0), width * 3, height,     cudaMemcpyDeviceToHost);
    LastError::setLastError("error: imageWriteJpegFromRgb cudaMemcpy2D");

    stbi_write_jpg(fileName.c_str(), width, height, 3, data, quality);

    free(data);

    return LastError::getLastErrorCode();
}

int imageWriteJpegFromBgr     (std::string fileName, const Image2D_Ptr& i2d, int quality)
{
    Image2D_Ptr i2dRgb = Image2D_Ptr(i2d.width(), i2d.height(), ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);
    int ret = bgrToRgb(i2d, i2dRgb);
    if (ret != 0)
        return ret;

    imageWriteJpegFromRgb(fileName, i2dRgb, quality);

    return LastError::getLastErrorCode();
}

int imageWriteJpegFromYuv420p (std::string fileName, const Image2D_Ptr& i2d, int quality)
{
    Image2D_Ptr i2dRgb = Image2D_Ptr(i2d.width(), i2d.height(), ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);
    int ret = yuv420pToRgb(i2d, i2dRgb);
    if (ret != 0)
        return ret;

    imageWriteJpegFromRgb(fileName, i2dRgb, quality);

    return LastError::getLastErrorCode();
}

int imageWriteJpegFromNv12    (std::string fileName, const Image2D_Ptr& i2d, int quality)
{
    Image2D_Ptr i2dRgb = Image2D_Ptr(i2d.width(), i2d.height(), ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);
    int ret = nv12ToRgb(i2d, i2dRgb);
    if (ret != 0)
        return ret;

    imageWriteJpegFromRgb(fileName, i2dRgb, quality);

    return LastError::getLastErrorCode();
}

int imageWriteBmpFromGray     (std::string fileName, const Image2D_Ptr& i2d)
{
    if((int)(i2d.type()) != (int)(ImageType::IMAGE_UNSIGNED_CHAR))
    {
        LastError::setLastError("error: Image2D_Ptr type must unsigned char", -1);
        return LastError::getLastErrorCode();
    }

    int width = i2d.width();
    int height = i2d.height();
    unsigned char* data = (unsigned char*)malloc(width * height);
    if (data == nullptr)
    {
        LastError::setLastError("error: imageWriteBmpFromGray malloc", -1);
        return LastError::getLastErrorCode();
    }

    cudaMemcpy2D(data, width, i2d.data(0), i2d.linesize(0), width ,height,     cudaMemcpyDeviceToHost);
    LastError::setLastError("error: imageWriteBmpFromGray cudaMemcpy2D");

    stbi_write_bmp(fileName.c_str(), width, height, 1, data);

    free(data);

    return LastError::getLastErrorCode();

}

int imageWriteBmpFromRgb      (std::string fileName, const Image2D_Ptr& i2d)
{
    if((int)(i2d.type()) != (int)(ImageType::IMAGE_UNSIGNED_CHAR))
    {
        LastError::setLastError("error: Image2D_Ptr type must unsigned char", -1);
        return LastError::getLastErrorCode();
    }

    int width = i2d.width();
    int height = i2d.height();
    unsigned char* data = (unsigned char*)malloc(width * height * 3);
    if (data == nullptr)
    {
        LastError::setLastError("error: imageWriteBmpFromRgb malloc", -1);
        return LastError::getLastErrorCode();
    }

    cudaMemcpy2D(data, width * 3, i2d.data(0), i2d.linesize(0), width * 3, height,     cudaMemcpyDeviceToHost);
    LastError::setLastError("error: imageWriteBmpFromRgb cudaMemcpy2D");

    stbi_write_bmp(fileName.c_str(), width, height, 3, data);

    free(data);

    return LastError::getLastErrorCode();
}

int imageWriteBmpFromBgr      (std::string fileName, const Image2D_Ptr& i2d)
{
    Image2D_Ptr i2dRgb = Image2D_Ptr(i2d.width(), i2d.height(), ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);
    int ret = bgrToRgb(i2d, i2dRgb);
    if (ret != 0)
        return ret;

    imageWriteBmpFromRgb(fileName, i2dRgb);

    return LastError::getLastErrorCode();
}

int imageWriteBmpFromYuv420p  (std::string fileName, const Image2D_Ptr& i2d)
{
    Image2D_Ptr i2dRgb = Image2D_Ptr(i2d.width(), i2d.height(), ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);
    int ret = yuv420pToRgb(i2d, i2dRgb);
    if (ret != 0)
        return ret;

    imageWriteBmpFromRgb(fileName, i2dRgb);

    return LastError::getLastErrorCode();
}

int imageWriteBmpFromNv12     (std::string fileName, const Image2D_Ptr& i2d)
{
    Image2D_Ptr i2dRgb = Image2D_Ptr(i2d.width(), i2d.height(), ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);
    int ret = nv12ToRgb(i2d, i2dRgb);
    if (ret != 0)
        return ret;

    imageWriteBmpFromRgb(fileName, i2dRgb);

    return LastError::getLastErrorCode();
}

int imageWritePngFromGray     (std::string fileName, const Image2D_Ptr& i2d)
{
    if((int)(i2d.type()) != (int)(ImageType::IMAGE_UNSIGNED_CHAR))
    {
        LastError::setLastError("error: Image2D_Ptr type must unsigned char", -1);
        return LastError::getLastErrorCode();
    }

    int width = i2d.width();
    int height = i2d.height();
    unsigned char* data = (unsigned char*)malloc(width * height);
    if (data == nullptr)
    {
        LastError::setLastError("error: imageWritePngFromGray malloc", -1);
        return LastError::getLastErrorCode();
    }

    cudaMemcpy2D(data, width, i2d.data(0), i2d.linesize(0), width ,height,     cudaMemcpyDeviceToHost);
    LastError::setLastError("error: imageWritePngFromGray cudaMemcpy2D");

    stbi_write_png(fileName.c_str(), width, height, 1, data, width);

    free(data);

    return LastError::getLastErrorCode();
}

int imageWritePngFromRgb      (std::string fileName, const Image2D_Ptr& i2d)
{
    if((int)(i2d.type()) != (int)(ImageType::IMAGE_UNSIGNED_CHAR))
    {
        LastError::setLastError("error: Image2D_Ptr type must unsigned char", -1);
        return LastError::getLastErrorCode();
    }

    int width = i2d.width();
    int height = i2d.height();
    unsigned char* data = (unsigned char*)malloc(width * height * 3);
    if (data == nullptr)
    {
        LastError::setLastError("error: imageWritePngFromRgb malloc", -1);
        return LastError::getLastErrorCode();
    }

    cudaMemcpy2D(data, width * 3, i2d.data(0), i2d.linesize(0), width * 3, height,     cudaMemcpyDeviceToHost);
    LastError::setLastError("error: imageWritePngFromRgb cudaMemcpy2D");

    stbi_write_png(fileName.c_str(), width, height, 3, data, width * 3);

    free(data);

    return LastError::getLastErrorCode();
}

int imageWritePngFromBgr      (std::string fileName, const Image2D_Ptr& i2d)
{
    Image2D_Ptr i2dRgb = Image2D_Ptr(i2d.width(), i2d.height(), ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);
    int ret = bgrToRgb(i2d, i2dRgb);
    if (ret != 0)
        return ret;

    imageWritePngFromRgb(fileName, i2dRgb);

    return LastError::getLastErrorCode();
}

int imageWritePngFromYuv420p  (std::string fileName, const Image2D_Ptr& i2d)
{
    Image2D_Ptr i2dRgb = Image2D_Ptr(i2d.width(), i2d.height(), ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);
    int ret = yuv420pToRgb(i2d, i2dRgb);
    if (ret != 0)
        return ret;

    imageWritePngFromRgb(fileName, i2dRgb);

    return LastError::getLastErrorCode();
}

int imageWritePngFromNv12     (std::string fileName, const Image2D_Ptr& i2d)
{
    Image2D_Ptr i2dRgb = Image2D_Ptr(i2d.width(), i2d.height(), ImageFormat::IMAGE_RGB, ImageType::IMAGE_UNSIGNED_CHAR);
    int ret = nv12ToRgb(i2d, i2dRgb);
    if (ret != 0)
        return ret;

    imageWritePngFromRgb(fileName, i2dRgb);

    return LastError::getLastErrorCode();
}

}